import base64
import io
import re

import requests
from PIL import ImageEnhance, Image
from pytesseract import pytesseract
from robot.api import logger

from Lib.zcm.exception.TestException import TestException
from Lib.zcm.exception.errorcode import ErrorCode

from Lib.zcm.assistant.FwVariables import get_value


class BaseOCR(object):
    def ocr_number_file(self, image_file, excepted_count: int = -1):
        # 识别验证码里只有数字的情况
        # excepted_count: 期望是多少位数字，-1的不管
        pass

    def ocr_general_file(self, image_file, excepted_count: int = -1):
        # 识别验证码包括英文等其它文字的情况
        # excepted_count: 期望是多少位，-1的不管
        pass

    def ocr_number_base64(self, base64, excepted_count: int = -1):
        pass

    def ocr_general_base64(self, base64, excepted_count: int = -1):
        pass

class BaiduOCR(BaseOCR):
    def __init__(self):
        self._client_id = get_value("OCR", "baidu_client_id", "RGf3RIgx6zIS5fzRrYFkwuFy")
        self._client_secret = get_value("OCR", "baidu_client_secret", "BGGO5GGYQxosU23yksxtCWkq9QiFXZkb")

    def _get_access_token(self):
        author_url = f'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={self._client_id}&client_secret={self._client_secret}'
        resp = requests.get(author_url, timeout=120)
        if resp:
            logger.debug(f"baidu OCR author response {resp}")
            resp_data = resp.json()
            return resp_data.get("access_token")
        return None

    def _request(self, url, access_token, params):
        request_url = f"{url}?access_token={access_token}"
        headers = {'content-type': 'application/x-www-form-urlencoded'}

        resp = requests.post(request_url, data=params, headers=headers, timeout=180)
        if resp:
            logger.debug(f"baidu OCR request response {resp}")
            resp_data = resp.json()
            error_code = resp_data.get("error_code")
            if error_code:
                error_message = resp_data.get("error_msg", "")
                raise TestException.by_error_code(ErrorCode.UNKNOWN, value=f"{error_message}(code={error_code})")

            words_result_num = resp_data.get("words_result_num", 0)
            if words_result_num <= 0:
                return ""   # 没有识别到内容

            words_result = resp_data.get("words_result", [])
            result = ""
            for word in words_result:
                result = result + word.get("words", "")
            return result

        raise TestException.by_error_code(ErrorCode.UNKNOWN, value=f"Baidu OCR request failed")

    def ocr_number_file(self, image_file, excepted_count: int = -1):
        # 二进制方式打开图片文件
        f = open(image_file, 'rb')
        img = base64.b64encode(f.read())
        return self.ocr_number_base64(img, excepted_count)

    def ocr_general_file(self, image_file, excepted_count: int = -1):
        f = open(image_file, 'rb')
        img = base64.b64encode(f.read())
        return self.ocr_general_base64(img, excepted_count)

    def ocr_number_base64(self, base64, excepted_count: int = -1):
        access_token = self._get_access_token()
        if not access_token:
            raise TestException.by_error_code(ErrorCode.UNKNOWN, value='Can not get Access Token when OCR image')

        request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/numbers"
        params = {"image": base64}

        result = self._request(request_url, access_token, params)

        if excepted_count != -1:
            if len(result) != excepted_count:
                logger.error(f"Baidu OCR failed, except count {excepted_count}, but result is {result}.")
                return result, False
        return result, True

    def ocr_general_base64(self, base64, excepted_count: int = -1):
        access_token = self._get_access_token()
        if not access_token:
            raise TestException.by_error_code(ErrorCode.UNKNOWN, value='Can not get Access Token when OCR image')

        request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/webimage"
        params = {"image": base64}

        result = self._request(request_url, access_token, params)

        if excepted_count != -1:
            if len(result) != excepted_count:
                logger.error(f"Baidu OCR failed, except count {excepted_count}, but result is {result}")
                return result, False
        return result, True

class ttshituOCR(BaseOCR):
    # 图
    def __init__(self):
        self._user_name = get_value("OCR", "ttshitu_user", "jameshuang")
        self._password = get_value("OCR", "ttshitu_password", "1122334455")
        self._api_url = get_value("OCR", "ttshitu_api_url", "http://api.ttshitu.com/predict")

    def _request(self, url, typeid, image):
        data = {"username": self._user_name,
                "password": self._password,
                "typeid": typeid,
                "image": image}
        resp = requests.post(url, json=data)
        if resp:
            logger.debug(f"ttshitu OCR request response {resp}")
            resp_data = resp.json()
            if not resp_data['success']:
                raise TestException.by_error_code(ErrorCode.UNKNOWN, value=f'{resp_data["message"]}')
            return resp_data["data"]["result"]
        raise TestException.by_error_code(ErrorCode.UNKNOWN, value=f"ttshitu OCR request failed")

    def ocr_number_base64(self, base64, excepted_count: int = -1):
        result = self._request(self._api_url, 1001, base64)

        if excepted_count != -1:
            if len(result) != excepted_count:
                logger.error(f"ttshitu OCR failed, except count {excepted_count}, but result is {result}")
                return result, False
        return result, True

    def ocr_general_base64(self, base64, excepted_count: int = -1):
        result = self._request(self._api_url, 1003, base64)

        if excepted_count != -1:
            if len(result) != excepted_count:
                logger.error(f"ttshitu OCR failed, except count {excepted_count}, but result is {result}")
                return result, False
        return result, True

    def ocr_number_file(self, image_file, excepted_count: int = -1):
        # 二进制方式打开图片文件
        f = open(image_file, 'rb')
        img = base64.b64encode(f.read())
        b64 = img.decode()
        return self.ocr_number_base64(b64, excepted_count)

    def ocr_general_file(self, image_file, excepted_count: int = -1):
        f = open(image_file, 'rb')
        img = base64.b64encode(f.read())
        b64 = img.decode()
        return self.ocr_general_base64(b64, excepted_count)

class LocalOCR(BaseOCR):
    def _recognition(self, image):
        image = image.convert('L')  # 转换模式：L | RGB
        image = ImageEnhance.Contrast(image)  # 增强对比度
        image = image.enhance(2.0)
        result = pytesseract.image_to_string(image)
        logger.debug(f"local ocr recognition result is {result}")
        return result

    def ocr_number_base64(self, base64_string, excepted_count: int = -1):
        data_base64 = base64.b64decode(base64_string)
        image_data = io.BytesIO(data_base64)
        image = Image.open(image_data)

        result = self._recognition(image)

        number = re.sub("\D", "", result) #提取出数字

        if excepted_count != -1:
            if len(number) != excepted_count:
                logger.error(f"local OCR failed, except count {excepted_count}, but result is {number}")
                return number, False
        return number, True

    def ocr_number_file(self, image_file, excepted_count: int = -1):

        image = Image.open(image_file)

        result = self._recognition(image)

        number = re.sub("\D", "", result)  # 提取出数字
        if excepted_count != -1:
            if len(number) != excepted_count:
                logger.error(f"local OCR failed, except count {excepted_count}, but result is {number}")
                return number, False
        return number, True

    def ocr_general_base64(self, base64_string, excepted_count: int = -1):
        data_base64 = base64.b64decode(base64_string)
        image_data = io.BytesIO(data_base64)
        image = Image.open(image_data)

        result = self._recognition(image)

        if excepted_count != -1:
            if len(result) != excepted_count:
                logger.error(f"local OCR failed, except count {excepted_count}, but result is {result}")
                return result, False
        return result, True

    def ocr_general_file(self, image_file, excepted_count: int = -1):
        image = Image.open(image_file)
        result = self._recognition(image)

        if excepted_count != -1:
            if len(result) != excepted_count:
                logger.error(f"local OCR failed, except count {excepted_count}, but result is {result}")
                return result, False
        return result, True

if __name__ == '__main__':
    ocr = LocalOCR()
    print(ocr.ocr_number_file(r"C:\Users\james\Downloads\ocr4.jpg", 4))

